from fastapi import FastAPI, HTTPException
from transformers import AutoModelForMaskedLM, AutoTokenizer, pipeline
import traceback
from pydantic import BaseModel

# 下载 CodeBERT 模型和分词器
MODEL_NAME = "microsoft/codebert-base"

# 加载模型和分词器
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForMaskedLM.from_pretrained(MODEL_NAME)

# 使用 Hugging Face Pipeline 搭建语言模型接口
code_completion_pipeline = pipeline("fill-mask", model=model, tokenizer=tokenizer)

# 创建 FastAPI 应用
app = FastAPI()


# 定义输入数据结构
class CodeInput(BaseModel):
    code_snippet: str  # 用户输入的代码片段，使用 <mask> 表示需要补全的部分


@app.post("/complete")
def complete_code(input_data: CodeInput):
    """
    输入代码片段，返回模型补全结果。
    示例输入： "int main() { return <mask>; }"
    """
    code_snippet = input_data.code_snippet
    print(f"Received code snippet: {code_snippet}")  # 打印接收到的代码片段

    # if "<mask>" not in code_snippet:
    #     raise HTTPException(
    #         status_code=400,
    #         detail="The code snippet must contain '<mask>' for completion.",
    #     )

    try:
        # 使用模型预测补全
        results = code_completion_pipeline(code_snippet)
        completions = [result["sequence"] for result in results]
        return {"input": code_snippet, "completions": completions}

    except Exception as e:
        # 捕获异常并输出错误信息
        print(f"Error occurred: {e}")
        traceback.print_exc()  # 打印详细的异常信息
        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")


@app.get("/")
def read_root():
    return {
        "message": "CodeBERT API is running. Use /complete endpoint for code completion."
    }
